#! /usr/bin/env python3
'''
author:toamin
function:代码由四部分组成,定位信息+规划算法+控制算法+Ros通信,实现单车点到点规划控制
'''

import sys
sys.path.append('/home/taomin/MVTP_ws_tm/src/one_car_planning_control/base/')
# from mvccep_image_localization.msg import base_localization # 获取定位信息
from geometry_msgs.msg import Twist
import rospy # ros发布订阅用到
import numpy as np # np数学库
from Parameter import Parameter as PA # A*栅格地图参数、车辆尺寸参数
import hybrid_astar as h_astar # 混合A*算法
import math
# import startP_sub as startP_sub
import matplotlib.pyplot as plt
import scipy.linalg as la # 线性代数求解库
import csv
from LQR import LQRControl
from pathPreProcess import pathPreProcess

# 车辆状态量
class State:
    def __init__(self, x, y, yaw, v):
        self.x = x
        self.y = y
        self.yaw = yaw # yaw单位为角度
        self.v = v # 小车当前速度

# 混合A*规划路径，输入：起点、终点（x, y, yaw）State类型； 输出：路径（list类型，每一个元素都是State类型）
def hybridAstarPathPlanning(startP, endP):
    """
    规划算法
    Args:startP (list)、endP (list)
    Returns:path 路径
    """
    # 起点终点预处理，转换为A*地图坐标 （实际场景中以矩形中点为坐标原点，A*地图中以左下角为坐标原点）
    startP[0] = startP[0] / PA.N + PA.MAP_SIZE[0] / 2
    startP[1] = startP[1] / PA.N + PA.MAP_SIZE[0] / 2
    endP[0] = endP[0] / PA.N + PA.MAP_SIZE[0] / 2
    endP[1] = endP[1] / PA.N + PA.MAP_SIZE[0] / 2
    
    # path是一个类实例，包含x, y, yaw, direction, cost属性，其中每个变量均为list
    path = h_astar.main(startP, endP) #调用混合A*进行路径规划
    
    # A*地图坐标转换为实际坐标（相对于预处理公式中的逆转换）
    for i in range(len(path.x)):
        path.x[i] = (path.x[i] - PA.MAP_SIZE[0]/2) * PA.N
        path.y[i] = (path.y[i] - PA.MAP_SIZE[1]/2) * PA.N
    
    return path

# def lqr_speed_steering_control(state, traj, pe, pth_e, Q, R):
#     ind, e = calc_nearest_index(state, traj.x, traj.y, traj.yaw) # e表示车辆离轨迹最近点的距离（也可以理解为位置误差），ind为该点的索引
#     tv = traj.v[ind]

#     k = traj.kappa[ind] # kappa为曲率 
#     v = state.v
#     th_e = pi_2_pi(state.yaw - traj.yaw[ind])

#     # A = [1.0, dt, 0.0, 0.0, 0.0
#     #      0.0, 0.0, v, 0.0, 0.0]
#     #      0.0, 0.0, 1.0, dt, 0.0]
#     #      0.0, 0.0, 0.0, 0.0, 0.0]
#     #      0.0, 0.0, 0.0, 0.0, 1.0]
#     A = np.zeros((5, 5))
#     A[0, 0] = 1.0
#     A[0, 1] = dt
#     A[1, 2] = v
#     A[2, 2] = 1.0
#     A[2, 3] = dt
#     A[4, 4] = 1.0

#     # B = [0.0, 0.0
#     #     0.0, 0.0
#     #     0.0, 0.0
#     #     v/L, 0.0
#     #     0.0, dt]
#     B = np.zeros((5, 2))
#     B[3, 0] = v / L
#     B[4, 1] = dt

#     K, _, _ = dlqr(A, B, Q, R)

#     # state vector
#     # x = [e, dot_e, th_e, dot_th_e, delta_v]
#     # e: lateral distance to the path
#     # dot_e: derivative of e
#     # th_e: angle difference to the path
#     # dot_th_e: derivative of th_e
#     # delta_v: difference between current speed and target speed
#     x = np.zeros((5, 1))
#     x[0, 0] = e
#     x[1, 0] = (e - pe) / dt
#     x[2, 0] = th_e
#     x[3, 0] = (th_e - pth_e) / dt
#     x[4, 0] = v - tv

#     # input vector
#     # u = [delta, accel]
#     # delta: steering angle
#     # accel: acceleration
#     ustar = -K @ x # @表示矩阵乘法 计算得到控制矩阵u

#     # calc steering input 基于跟踪误差模型（考虑道路弯曲）龚建伟MPC/p30
#     ff = math.atan2(L * k, 1)  # feedforward steering angle  航向误差需要的偏转角
#     fb = pi_2_pi(ustar[0, 0])  # feedback steering angle   横向误差需要的偏转角
#     delta = ff + fb

#     if traj.direction[ind] < 0:
#         delta *= -1

#     # calc accel input
#     accel = ustar[1, 0]

#     return delta, ind, accel, e, th_e

# def calc_nearest_index(state, cx, cy, cyaw):
#     dx = [state.x - icx for icx in cx]
#     dy = [state.y - icy for icy in cy]

#     d = [idx ** 2 + idy ** 2 for (idx, idy) in zip(dx, dy)]

#     min_d = min(d)

#     ind = d.index(min_d)

#     min_d = math.sqrt(min_d)

#     dxl = cx[ind] - state.x
#     dyl = cy[ind] - state.y

#     angle = pi_2_pi(cyaw[ind] - math.atan2(dyl, dxl))
#     if angle < 0: # 小于0，车辆在轨迹的右侧，大于0，车辆在轨迹的左侧
#         min_d *= -1

#     return ind, min_d

# def pi_2_pi(angle):
#     return (angle + math.pi) % (2 * math.pi) - math.pi

# def dlqr(A, B, Q, R):
#     """Solve the discrete time lqr controller.
#     x[k+1] = A x[k] + B u[k]
#     cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]
#     # ref Bertsekas, p.151
#     """

#     # first, try to solve the ricatti equation
#     P = solve_dare(A, B, Q, R)

#     # compute the LQR gain
#     K = la.inv(B.T @ P @ B + R) @ (B.T @ P @ A)

#     eig_result = la.eig(A - B @ K)

#     return K, P, eig_result[0]

# def solve_dare(A, B, Q, R):
    """
    solve a discrete time_Algebraic Riccati equation (DARE)
    """
    p = Q
    p_next = Q
    max_iter = 150
    eps = 0.01

    for i in range(max_iter):
        p_next = A.T @ p @ A - A.T @ p @ B @ \
                 la.inv(R + B.T @ p @ B) @ B.T @ p @ A + Q
        if (abs(p_next - p)).max() < eps:
            break
        p = p_next

    return p_next

class OneCarControl(object):
    """
    path已知,由控制算法得到车辆控制量,随后Ros进行发布
    """
    def __init__(self, trajectory, traj_list, robot_id):
        self.robot_id = robot_id
        self.trajectory = trajectory
        self.traj_list = traj_list # 将轨迹分段进行控制
        # 定义订阅者，订阅位置信息，Subscribe类的实例化（话题名， 消息类型， 回调函数）
        self.draw_state = State([],[],[],[])
        self.sub = rospy.Subscriber("base_localization", base_localization, self.controlCallback, queue_size = 30)
        # 定义发布者
        self.pub = rospy.Publisher("/robot_{}/cmd_vel".format(robot_id), Twist, queue_size = 30)
        self.lqrcontroller = LQRControl()
        
    # 回调函数是系统进行传参的，以下msg接收的值是话题/base_localization发布出来的消息
    # (相关命令 rostopic echo /base_localization)
    def controlCallback(self, msg):
        # 从定位拿到小车实时状态
        self.robot_list = [[msg.robot_0.pos.x, msg.robot_0.pos.y, msg.robot_0.angle, msg.robot_0.twist_linear.x], 
                           [msg.robot_1.pos.x, msg.robot_1.pos.y, msg.robot_1.angle, msg.robot_1.twist_linear.x],
                           [msg.robot_2.pos.x, msg.robot_2.pos.y, msg.robot_2.angle, msg.robot_2.twist_linear.x],
                           [msg.robot_3.pos.x, msg.robot_3.pos.y, msg.robot_3.angle, msg.robot_3.twist_linear.x],
                           [msg.robot_4.pos.x, msg.robot_4.pos.y, msg.robot_4.angle, msg.robot_4.twist_linear.x]]
        self.state = State(self.robot_list[self.robot_id][0], self.robot_list[self.robot_id][1], 
                           self.robot_list[self.robot_id][2], self.robot_list[self.robot_id][3])
        
        self.draw_state.x.append(self.state.x)
        self.draw_state.y.append(self.state.y)
        self.draw_state.yaw.append(self.state.yaw)
        self.draw_state.v.append(self.state.v) # 画图用

        for i in range(len(self.traj_list)): #xxxxxxxxxxxxxxxxxxxxxxxxx
            if not self.lqrcontroller.STOP_FLAG:
                # 纯跟踪算法，传入小车实时状态、两个数组(即一系列路径点的二维坐标)，得到每个点对应的方向盘转角
                vel_msg = Twist() # 定义一条Twist类型的消息
                # vel_msg.linear.x, vel_msg.angular.z = self.lqrcontroller.lqr_speed_steering_control(self.state, self.traj_list[i]) # 分段控制
                vel_msg.linear.x, vel_msg.angular.z = self.lqrcontroller.lqr_speed_steering_control(self.state, self.trajectory) # 整段控制     
                # 发布线速度加速度信息
                self.pub.publish(vel_msg)
                rospy.loginfo("Publish velocity command [%.2f m/s, %.2f rad/s]", vel_msg.linear.x, vel_msg.angular.z)
                                
            else:
                # 自动结束ros发布的命令
                rospy.loginfo("*** Arrived the end point! ---------->>")
                rospy.signal_shutdown('Over')
            
def drawTraj(draw_state, traj):
    # 设置横纵坐标的名称以及对应字体格式
    font2 = {'family' : 'Times New Roman', 'weight' : 'normal', 'size' : 12}
    plt.title("轨迹对比图")
    plt.plot(draw_state.x, draw_state.y, color = "r", marker = ".", label = "真实轨迹")
    plt.plot(traj.x, traj.y, color = "b", linewidth = 2, label = "参考轨迹")
                
    plt.xlabel('x(m)',font2)
    plt.ylabel('y(m)',font2)
    plt.grid() # 显示网格
    plt.legend(loc = 'upper right', fontsize=12)  # 将图例显示在右上方
    plt.axis("equal") # 将 x 轴和 y 轴的比例设置为相等
    
    plt.show()

def getParam():
    try:
        rospy.get_param('~Flag')
        robot_id = rospy.get_param('~robot_id')
        endP = rospy.get_param('~endP') # 运行的时候通过命令行给定
        endP = endP.split(",")
        endP = [float(endP[0]), float(endP[1]), round(math.radians(float(endP[2])), 2)]
        # 设置endP格式:输入[0.52,0.52,90] -> 转化为[0.52,0.52,1.57] (转成弧度并保留两位小数)
    except KeyError:
        robot_id = 3
        # endP = [0.6, -0.52, math.radians(-90)] # 调试的时候给固定值
        angle = 90
        endP = [0.52, -0.52, math.radians(angle)] # 调试的时候给固定值
    except Exception:
        print("请正确输入参数，例如:roslaunch one_car_planning_control one_car_planning_control endP:=0,-0.52,90")
    return robot_id, endP

def drawPath(path, Point_O):
    # plt.title("轨迹对比图")
    # plt.plot(draw_state.x, draw_state.y, color = "r", marker = ".", label = "真实轨迹")
    # plt.plot(path.x, path.y, color = "b", linewidth = 2, label = "参考轨迹")
    plt.scatter(path.x, path.y, s=4, color = "b")

    for i in range(len(Point_O)):
        plt.scatter(Point_O[i][0], Point_O[i][1], s=18, color='g') # 画出尖点
    
    direct = path.direction[0]
    for i in range(len(path.x)):
        if path.direction[i] is not direct:
            direct = path.direction[i]
            plt.scatter(path.x[i - 1], path.y[i - 1], s=18, color='r')
    
    # plt.plot(draw_state.x, draw_state.v, color = "green", label = "速度")
    # plt.plot(path.x, path.direction, color = "y", label = "Direction")
    # 设置横纵坐标的名称以及对应字体格式
    # font2 = {'family' : 'Times New Roman', 'weight' : 'normal', 'size' : 12}
    # plt.xlabel('x(m)',font2)
    # plt.ylabel('y(m)',font2)
    plt.grid() # 显示网格
    # plt.legend(loc = 'upper right', fontsize=12)  # 将图例显示在右上方
    plt.axis("equal") # 将 x 轴和 y 轴的比例设置为相等
    plt.show()

if __name__ == '__main__':
    try:
        print("*** 开始执行主程序--------------->>\n")
        # 初始化一个节点，anonymous表示是否在节点名后追加随机数以确保节点名的唯一性
        rospy.init_node("planning_control_node", anonymous=False)
        # 获取参数
        # robot_id, endP = getParam() # 运行的时候通过命令行给定
        # startP = startP_sub.startP_sub(robot_id)# 从定位获取车辆的起始点（当前点）
        # print("*** 被控车辆: ", robot_id, "号")
        # print("*** 车辆起点: ", startP)
        # print("*** 车辆终点: ", endP, "\n")
        startP = [0,0,0]
        endP = [0.52,-0.6,90]
        # startP = [0,0,0]
        # endP = [0,-0.6,0]
        # A*规划路径
        path = hybridAstarPathPlanning(startP, endP) # 调用混合A*算法，传入起点与终点后规划出轨迹并return
        print(len(path.x), "  ", len(path.y), "  ", len(path.yaw), "  ", len(path.direction),"\n")
        # 路径预处理（为解决倒车问题）
        trajectory, traj_list, Point_O= pathPreProcess(path)   ## kappa计算还存在问题（拿文件储存路径信息，比较各种曲率计算方法）
        
        # 构建控制器
        # trajController = OneCarControl(trajectory, traj_list, robot_id)
        # rospy.spin()
        # 轨迹对比图
        # drawTraj(trajController.draw_state, trajectory)

        with open('output.csv', 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            for traj_x, traj_y in zip(trajectory.x, trajectory.y):
                writer.writerow([traj_x, traj_y]) # 将轨迹坐标离散点写入文件

        print(len(trajectory.x), " ", len(trajectory.y), " ", len(trajectory.yaw), " ", len(trajectory.direction), " ", len(trajectory.kappa), "\n")

        for i in range(len(trajectory.x)):
            print("ind = ", i, "  direction = ", trajectory.direction[i],"  yaw = ", math.degrees(trajectory.yaw[i]), "  kappa = ", trajectory.kappa[i])
        drawPath(trajectory, Point_O) # 画散点图
        # drawPath(path)
        
    except rospy.ROSInterruptException:
        rospy.signal_shutdown("err")
